package com.cyx.exercise.ollama;

import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Flow;
import java.util.function.Consumer;

/**
 * Ollama API客户端 - 使用Java原生HTTP客户端
 * 支持同步和异步调用，以及流式响应处理
 */
public class OllamaClient {
    
    private final HttpClient httpClient;
    private final String baseUrl;
    
    public OllamaClient(String baseUrl) {
        this.baseUrl = baseUrl.endsWith("/") ? baseUrl.substring(0, baseUrl.length() - 1) : baseUrl;
        this.httpClient = HttpClient.newBuilder()
                .connectTimeout(Duration.ofSeconds(30))
                .build();
    }
    
    /**
     * 同步调用Ollama生成接口
     */
    public OllamaResponse generate(OllamaRequest request) throws IOException, InterruptedException {
        String jsonBody = buildJsonRequest(request);
        
        HttpRequest httpRequest = HttpRequest.newBuilder()
                .uri(URI.create(baseUrl + "/api/generate"))
                .header("Content-Type", "application/json")
                .header("Accept", "application/json")
                .timeout(Duration.ofMinutes(5))
                .POST(HttpRequest.BodyPublishers.ofString(jsonBody))
                .build();
        
        HttpResponse<String> response = httpClient.send(httpRequest, HttpResponse.BodyHandlers.ofString());
        
        if (response.statusCode() != 200) {
            throw new IOException("HTTP Error: " + response.statusCode() + " - " + response.body());
        }
        
        return parseResponse(response.body());
    }
    
    /**
     * 异步调用Ollama生成接口
     */
    public CompletableFuture<OllamaResponse> generateAsync(OllamaRequest request) {
        String jsonBody = buildJsonRequest(request);
        
        HttpRequest httpRequest = HttpRequest.newBuilder()
                .uri(URI.create(baseUrl + "/api/generate"))
                .header("Content-Type", "application/json")
                .header("Accept", "application/json")
                .timeout(Duration.ofMinutes(5))
                .POST(HttpRequest.BodyPublishers.ofString(jsonBody))
                .build();
        
        return httpClient.sendAsync(httpRequest, HttpResponse.BodyHandlers.ofString())
                .thenApply(response -> {
                    if (response.statusCode() != 200) {
                        throw new RuntimeException("HTTP Error: " + response.statusCode() + " - " + response.body());
                    }
                    return parseResponse(response.body());
                });
    }
    
    /**
     * 流式调用Ollama生成接口
     */
    public void generateStream(OllamaRequest request, Consumer<String> onChunk, Consumer<Exception> onError) {
        OllamaRequest streamRequest = new OllamaRequest(request.getModel(), request.getPrompt(), true);
        String jsonBody = buildJsonRequest(streamRequest);
        
        HttpRequest httpRequest = HttpRequest.newBuilder()
                .uri(URI.create(baseUrl + "/api/generate"))
                .header("Content-Type", "application/json")
                .header("Accept", "application/json")
                .timeout(Duration.ofMinutes(5))
                .POST(HttpRequest.BodyPublishers.ofString(jsonBody))
                .build();
        
        httpClient.sendAsync(httpRequest, HttpResponse.BodyHandlers.ofLines())
                .thenAccept(response -> {
                    if (response.statusCode() != 200) {
                        onError.accept(new IOException("HTTP Error: " + response.statusCode()));
                        return;
                    }
                    
                    response.body().forEach(line -> {
                        if (!line.trim().isEmpty()) {
                            try {
                                OllamaStreamResponse streamResponse = parseStreamResponse(line);
                                if (streamResponse.getResponse() != null) {
                                    onChunk.accept(streamResponse.getResponse());
                                }
                            } catch (Exception e) {
                                onError.accept(e);
                            }
                        }
                    });
                })
                .exceptionally(throwable -> {
                    onError.accept(new RuntimeException(throwable));
                    return null;
                });
    }
    
    /**
     * 调用Ollama聊天接口
     */
    public OllamaChatResponse chat(OllamaChatRequest request) throws IOException, InterruptedException {
        String jsonBody = buildChatJsonRequest(request);
        
        HttpRequest httpRequest = HttpRequest.newBuilder()
                .uri(URI.create(baseUrl + "/api/chat"))
                .header("Content-Type", "application/json")
                .header("Accept", "application/json")
                .timeout(Duration.ofMinutes(5))
                .POST(HttpRequest.BodyPublishers.ofString(jsonBody))
                .build();
        
        HttpResponse<String> response = httpClient.send(httpRequest, HttpResponse.BodyHandlers.ofString());
        
        if (response.statusCode() != 200) {
            throw new IOException("HTTP Error: " + response.statusCode() + " - " + response.body());
        }
        
        return parseChatResponse(response.body());
    }
    
    /**
     * 获取可用模型列表
     */
    public String getModels() throws IOException, InterruptedException {
        HttpRequest httpRequest = HttpRequest.newBuilder()
                .uri(URI.create(baseUrl + "/api/tags"))
                .header("Accept", "application/json")
                .timeout(Duration.ofSeconds(30))
                .GET()
                .build();
        
        HttpResponse<String> response = httpClient.send(httpRequest, HttpResponse.BodyHandlers.ofString());
        
        if (response.statusCode() != 200) {
            throw new IOException("HTTP Error: " + response.statusCode() + " - " + response.body());
        }
        
        return response.body();
    }
    
    // 构建JSON请求体
    private String buildJsonRequest(OllamaRequest request) {
        StringBuilder json = new StringBuilder();
        json.append("{");
        json.append("\"model\":\"").append(escapeJson(request.getModel())).append("\",");
        json.append("\"prompt\":\"").append(escapeJson(request.getPrompt())).append("\"");
        // 显式指定是否流式，避免默认开启流模式导致解析只取到第一块
        json.append(",\"stream\":").append(request.isStream() ? "true" : "false");
        json.append("}");
        return json.toString();
    }
    
    // 构建聊天JSON请求体
    private String buildChatJsonRequest(OllamaChatRequest request) {
        StringBuilder json = new StringBuilder();
        json.append("{");
        json.append("\"model\":\"").append(escapeJson(request.getModel())).append("\",");
        json.append("\"messages\":[");
        for (int i = 0; i < request.getMessages().size(); i++) {
            if (i > 0) json.append(",");
            OllamaChatMessage msg = request.getMessages().get(i);
            json.append("{\"role\":\"").append(escapeJson(msg.getRole())).append("\",");
            json.append("\"content\":\"").append(escapeJson(msg.getContent())).append("\"}");
        }
        json.append("]");
        // 显式指定是否流式
        json.append(",\"stream\":").append(request.isStream() ? "true" : "false");
        json.append("}");
        return json.toString();
    }
    
    // 简单的JSON转义
    private String escapeJson(String str) {
        if (str == null) return "";
        return str.replace("\\", "\\\\")
                  .replace("\"", "\\\"")
                  .replace("\n", "\\n")
                  .replace("\r", "\\r")
                  .replace("\t", "\\t");
    }
    
    // 解析响应（同时兼容非流式和误返回的流式多行JSON）
    private OllamaResponse parseResponse(String jsonResponse) {
        if (jsonResponse == null) return new OllamaResponse("", true);
        // 如果是多行JSON（流式），逐行累加response并以最后一个done为准
        String[] lines = jsonResponse.split("\r?\n");
        if (lines.length > 1) {
            StringBuilder sb = new StringBuilder();
            boolean done = false;
            for (String line : lines) {
                if (line == null || line.trim().isEmpty()) continue;
                String chunk = extractJsonValue(line, "response");
                if (chunk != null) sb.append(chunk);
                String doneVal = extractJsonValue(line, "done");
                if ("true".equalsIgnoreCase(doneVal)) done = true;
            }
            return new OllamaResponse(sb.toString(), done);
        }
        // 单行正常JSON
        String response = extractJsonValue(jsonResponse, "response");
        boolean done = "true".equalsIgnoreCase(extractJsonValue(jsonResponse, "done"));
        return new OllamaResponse(response, done);
    }
    
    // 解析流式响应
    private OllamaStreamResponse parseStreamResponse(String jsonLine) {
        String response = extractJsonValue(jsonLine, "response");
        boolean done = "true".equals(extractJsonValue(jsonLine, "done"));
        return new OllamaStreamResponse(response, done);
    }
    
    // 解析聊天响应
    private OllamaChatResponse parseChatResponse(String jsonResponse) {
        // 提取message对象中的content
        String messageStart = "\"message\":{";
        int messageIndex = jsonResponse.indexOf(messageStart);
        if (messageIndex != -1) {
            String messageJson = jsonResponse.substring(messageIndex + messageStart.length());
            int messageEnd = findMatchingBrace(messageJson);
            if (messageEnd != -1) {
                String messageContent = messageJson.substring(0, messageEnd);
                String content = extractJsonValue(messageContent, "content");
                String role = extractJsonValue(messageContent, "role");
                boolean done = "true".equals(extractJsonValue(jsonResponse, "done"));
                return new OllamaChatResponse(new OllamaChatMessage(role, content), done);
            }
        }
        return new OllamaChatResponse(new OllamaChatMessage("assistant", ""), true);
    }
    
    // 改进的JSON值提取
    private String extractJsonValue(String json, String key) {
        String searchKey = "\"" + key + "\":\"";
        int startIndex = json.indexOf(searchKey);
        if (startIndex == -1) {
            // 尝试布尔值或数字
            searchKey = "\"" + key + "\":";
            startIndex = json.indexOf(searchKey);
            if (startIndex == -1) return null;
            startIndex += searchKey.length();
            int endIndex = json.indexOf(",", startIndex);
            if (endIndex == -1) endIndex = json.indexOf("}", startIndex);
            if (endIndex == -1) return null;
            String value = json.substring(startIndex, endIndex).trim();
            // 移除可能的引号
            if (value.startsWith("\"") && value.endsWith("\"")) {
                value = value.substring(1, value.length() - 1);
            }
            return value;
        }

        startIndex += searchKey.length();
        int endIndex = startIndex;
        int braceCount = 0;
        boolean inString = true; // 我们已经在字符串内了

        while (endIndex < json.length()) {
            char c = json.charAt(endIndex);

            if (c == '\\' && inString) {
                endIndex += 2; // 跳过转义字符
                continue;
            }

            if (c == '"' && inString) {
                // 找到字符串结束
                break;
            }

            endIndex++;
        }

        if (endIndex >= json.length()) return null;
        String result = json.substring(startIndex, endIndex);

        // 处理Unicode转义序列
        result = unescapeJson(result);

        return result;
    }

    // 处理JSON转义序列
    private String unescapeJson(String str) {
        if (str == null) return null;
        return str.replace("\\\"", "\"")
                  .replace("\\\\", "\\")
                  .replace("\\n", "\n")
                  .replace("\\r", "\r")
                  .replace("\\t", "\t")
                  .replace("\\u003c", "<")
                  .replace("\\u003e", ">")
                  .replace("\\u0026", "&");
    }
    
    // 找到匹配的大括号
    private int findMatchingBrace(String json) {
        int braceCount = 1;
        for (int i = 0; i < json.length(); i++) {
            char c = json.charAt(i);
            if (c == '{') braceCount++;
            else if (c == '}') {
                braceCount--;
                if (braceCount == 0) return i;
            }
        }
        return -1;
    }
    
    public void close() {
        // HttpClient会自动管理连接池，无需显式关闭
    }
}
